package com.alibaba.datax.core.transport.transformer;

import com.alibaba.datax.common.element.Column;
import com.alibaba.datax.common.element.Record;
import com.alibaba.datax.common.element.StringColumn;
import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.transformer.Transformer;
import org.apache.commons.lang3.time.DateFormatUtils;

import java.util.Arrays;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author huang lei
 * @date 2022/1/20 10:58
 */
public class BatchNumberTransformer extends Transformer {

    private Map<String, AtomicInteger> mapCount = new ConcurrentHashMap<>();

    public BatchNumberTransformer() {
        setTransformerName("dx_number");
    }

    @Override
    public Record evaluate(Record record, Object... paras) {
        int columnIndex;
        int size;
        int srcColumnIndex = -1;
        try {
            if (paras.length < 2) {
                throw new RuntimeException("dx_reverse paras must gte 2");
            }

            columnIndex = (Integer) paras[0];
            size = Integer.parseInt((String) paras[1]);
            if (paras.length == 3) {
                srcColumnIndex = Integer.parseInt((String) paras[2]);
            }
        } catch (Exception e) {
            throw DataXException.asDataXException(TransformerErrorCode.TRANSFORMER_ILLEGAL_PARAMETER, "paras:" + Arrays.asList(paras).toString() + " => " + e.getMessage());
        }

        try {
            String prefix = "";
            if (srcColumnIndex >= 0) {
                Column srcColumn = record.getColumn(srcColumnIndex);
                Column.Type type = srcColumn.getType();
                if (type == Column.Type.DATE) {
                    prefix = DateFormatUtils.format(srcColumn.asDate(), "yyMMdd");
                } else {
                    prefix = srcColumn.asString();
                }
            } else if (srcColumnIndex == -1) {
                prefix = DateFormatUtils.format(new Date(), "yyMMdd");
            }

            AtomicInteger atomicInteger = mapCount.computeIfAbsent(prefix, s -> new AtomicInteger());
            int number = atomicInteger.getAndIncrement();
            int i = number / size;

            record.setColumn(columnIndex, new StringColumn(prefix + i));

        } catch (Exception e) {
            throw DataXException.asDataXException(TransformerErrorCode.TRANSFORMER_RUN_EXCEPTION, e.getMessage(), e);
        }
        return record;

    }
}
